{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ea51cf49-453a-4740-9fea-8af40b5d4909",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --quiet langchain langchain-community langchain-anthropic neo4j"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "98f1dbf6-facc-4be9-b4a1-30aaa0a768ba",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/tomazbratanic/anaconda3/lib/python3.11/site-packages/pandas/core/arrays/masked.py:60: UserWarning: Pandas requires version '1.3.6' or newer of 'bottleneck' (version '1.3.5' currently installed).\n",
      "  from pandas.core import (\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>question</th>\n",
       "      <th>type</th>\n",
       "      <th>database</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>What are the top 5 movies with a runtime great...</td>\n",
       "      <td>Simple Retrieval Queries</td>\n",
       "      <td>recommendations</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>List the first 3 directors born before 1950.</td>\n",
       "      <td>Simple Retrieval Queries</td>\n",
       "      <td>recommendations</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Which 5 users have rated more than 20 movies?</td>\n",
       "      <td>Simple Retrieval Queries</td>\n",
       "      <td>recommendations</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Identify the top 5 actors who have acted in mo...</td>\n",
       "      <td>Simple Retrieval Queries</td>\n",
       "      <td>recommendations</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>What are the top 3 genres associated with movi...</td>\n",
       "      <td>Simple Retrieval Queries</td>\n",
       "      <td>recommendations</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                            question  \\\n",
       "0  What are the top 5 movies with a runtime great...   \n",
       "1       List the first 3 directors born before 1950.   \n",
       "2      Which 5 users have rated more than 20 movies?   \n",
       "3  Identify the top 5 actors who have acted in mo...   \n",
       "4  What are the top 3 genres associated with movi...   \n",
       "\n",
       "                       type         database  \n",
       "0  Simple Retrieval Queries  recommendations  \n",
       "1  Simple Retrieval Queries  recommendations  \n",
       "2  Simple Retrieval Queries  recommendations  \n",
       "3  Simple Retrieval Queries  recommendations  \n",
       "4  Simple Retrieval Queries  recommendations  "
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "questions = pd.read_csv(\"text2cypher_questions.csv\")\n",
    "questions.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "df726b72-5ad2-41a5-b092-37851d36863c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4991\n"
     ]
    }
   ],
   "source": [
    "selected_questions = questions[\n",
    "    questions[\"database\"].isin(\n",
    "        [\n",
    "            \"recommendations\",\n",
    "            \"companies\",\n",
    "            \"twitch\",\n",
    "            \"twitter\",\n",
    "            \"gameofthrones\",\n",
    "            \"movies\",\n",
    "            \"neoflix\",\n",
    "        ]\n",
    "    )\n",
    "]\n",
    "print(len(selected_questions))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "97668508-3d62-4ba2-bb0d-c5be2f2d50f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "schemas = pd.read_csv('text2cypher_schemas.csv')\n",
    "schemas.head()\n",
    "schema_dict = {}\n",
    "for i, row in schemas.iterrows():\n",
    "    schema_dict[row['database']] = row['schema']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a4699c89-66a2-4f79-804f-e488efd4a744",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List, Union\n",
    "\n",
    "from langchain.chains.graph_qa.cypher_utils import CypherQueryCorrector, Schema\n",
    "from langchain_community.graphs import Neo4jGraph\n",
    "from langchain_core.messages import (\n",
    "    AIMessage,\n",
    "    SystemMessage,\n",
    "    ToolMessage,\n",
    ")\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import (\n",
    "    ChatPromptTemplate,\n",
    "    HumanMessagePromptTemplate,\n",
    "    MessagesPlaceholder,\n",
    ")\n",
    "from langchain_core.pydantic_v1 import BaseModel\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "from langchain_anthropic import ChatAnthropic\n",
    "import os\n",
    "\n",
    "os.environ[\"ANTHROPIC_API_KEY\"] = 'sk-'\n",
    "# LLMs\n",
    "cypher_llm = ChatAnthropic(model=\"claude-3-opus-20240229\", temperature=0.0, default_request_timeout=20)\n",
    "\n",
    "# Generate Cypher statement based on natural language input\n",
    "cypher_template = \"\"\"Based on the Neo4j graph schema below, write a Cypher query that would answer the user's question:\n",
    "{schema}\n",
    "\n",
    "Question: {question}\n",
    "Cypher query:\"\"\"  # noqa: E501\n",
    "\n",
    "cypher_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"\"\"Given an input question, convert it to a Cypher query. No pre-amble.\n",
    "Additional instructions:\n",
    "- Ensure that queries checking for non-null properties use `IS NOT NULL` in a straightforward manner.\n",
    "- Don't use `size((n)--(m))` for counting relationships. Instead use the new `count{{(n)--(m)}}` syntax.\n",
    "- Incorporate the new existential subqueries in examples where the query needs to check for the existence of a pattern.\n",
    "  Example: MATCH (p:Person)-[r:IS_FRIENDS_WITH]->(friend:Person)\n",
    "            WHERE exists{{ (p)-[:WORKS_FOR]->(:Company {{name: 'Neo4j'}})}}\n",
    "            RETURN p, r, friend\"\"\"        ),\n",
    "        (\"human\", cypher_template),\n",
    "    ]\n",
    ")\n",
    "\n",
    "cypher_response = (\n",
    "    cypher_prompt\n",
    "    | cypher_llm\n",
    "    | StrOutputParser()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c22eee16-8b01-4378-a4a5-5a6217720f16",
   "metadata": {},
   "outputs": [],
   "source": [
    "cypher_responses = []\n",
    "for i, row in selected_questions.iterrows():\n",
    "    print(i)\n",
    "    schema = schema_dict[row['database']]\n",
    "    try:\n",
    "        cypher = cypher_response.invoke({\"question\": row[\"question\"], \"schema\": schema})\n",
    "        cypher_responses.append({\"question\": row[\"question\"], \"database\": row[\"database\"], \"cypher\":cypher, \"type\":row[\"type\"]})\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        #print('timeout' + str(e))\n",
    "        #cypher_responses.append(\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "2e7c7239-8b53-4b78-a7fe-752a9f2ffae5",
   "metadata": {},
   "outputs": [],
   "source": [
    "combined_df = pd.DataFrame().from_records(cypher_responses)\n",
    "DEMO_URL = \"neo4j+s://demo.neo4jlabs.com\" "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87d899ed-7657-4526-936a-71e97155ea4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "syntax_error = []\n",
    "returns_results = []\n",
    "timeouts = []\n",
    "last_graph = \"\"\n",
    "for i, row in combined_df.reset_index().iterrows():\n",
    "    print(i)\n",
    "    if i % 100 == 0:\n",
    "        print(i)\n",
    "    # To avoid a new driver for every request\n",
    "    if row['database'] != last_graph:\n",
    "        last_graph = row[\"database\"]\n",
    "        print(last_graph)\n",
    "        graph = Neo4jGraph(\n",
    "            url=DEMO_URL, \n",
    "            username=row[\"database\"], \n",
    "            password=row[\"database\"], \n",
    "            database=row[\"database\"], \n",
    "            refresh_schema=False, \n",
    "            timeout=10)\n",
    "    try:\n",
    "        data = graph.query(row['cypher'])\n",
    "        if data:\n",
    "            returns_results.append(True)\n",
    "        else:\n",
    "            returns_results.append(False)\n",
    "        syntax_error.append(False)\n",
    "        timeouts.append(False)\n",
    "    except ValueError as e:\n",
    "        if \"Generated Cypher Statement is not valid\" in str(e):\n",
    "            syntax_error.append(True)\n",
    "            print(f\"Syntax error in Cypher query: {e}\")\n",
    "        else:\n",
    "            syntax_error.append(False)\n",
    "            print(f\"Other ValueError: {e}\")\n",
    "        returns_results.append(False)\n",
    "        timeouts.append(False)\n",
    "    except Exception as e:\n",
    "        if e.code  == \"Neo.ClientError.Transaction.TransactionTimedOutClientConfiguration\":\n",
    "            returns_results.append(False)\n",
    "            syntax_error.append(False)\n",
    "            timeouts.append(True)\n",
    "        else:\n",
    "            returns_results.append(False)\n",
    "            syntax_error.append(False)\n",
    "            timeouts.append(True)\n",
    "            \n",
    "        \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "8a0276a4-56e2-4955-8b27-d6fc5e8464d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "combined_df[\"syntax_error\"] = syntax_error\n",
    "combined_df[\"timeout\"] = timeouts\n",
    "combined_df[\"returns_results\"] = returns_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "d7056cd2-f099-41bf-a2dc-1e7badb82743",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>question</th>\n",
       "      <th>cypher</th>\n",
       "      <th>type</th>\n",
       "      <th>database</th>\n",
       "      <th>syntax_error</th>\n",
       "      <th>timeout</th>\n",
       "      <th>returns_results</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>What are the top 5 movies with a runtime great...</td>\n",
       "      <td>MATCH (m:Movie)\\nWHERE m.runtime &gt; 120\\nRETURN...</td>\n",
       "      <td>Simple Retrieval Queries</td>\n",
       "      <td>recommendations</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>List the first 3 directors born before 1950.</td>\n",
       "      <td>MATCH (d:Director)\\nWHERE d.born &lt; date('1950-...</td>\n",
       "      <td>Simple Retrieval Queries</td>\n",
       "      <td>recommendations</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Which 5 users have rated more than 20 movies?</td>\n",
       "      <td>MATCH (u:User)-[r:RATED]-&gt;(m:Movie)\\nWITH u, c...</td>\n",
       "      <td>Simple Retrieval Queries</td>\n",
       "      <td>recommendations</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Identify the top 5 actors who have acted in mo...</td>\n",
       "      <td>MATCH (a:Actor)-[:ACTED_IN]-&gt;(m:Movie)\\nWITH a...</td>\n",
       "      <td>Simple Retrieval Queries</td>\n",
       "      <td>recommendations</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>What are the top 3 genres associated with movi...</td>\n",
       "      <td>MATCH (m:Movie)-[:IN_GENRE]-&gt;(g:Genre)\\nWHERE ...</td>\n",
       "      <td>Simple Retrieval Queries</td>\n",
       "      <td>recommendations</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                            question  \\\n",
       "0  What are the top 5 movies with a runtime great...   \n",
       "1       List the first 3 directors born before 1950.   \n",
       "2      Which 5 users have rated more than 20 movies?   \n",
       "3  Identify the top 5 actors who have acted in mo...   \n",
       "4  What are the top 3 genres associated with movi...   \n",
       "\n",
       "                                              cypher  \\\n",
       "0  MATCH (m:Movie)\\nWHERE m.runtime > 120\\nRETURN...   \n",
       "1  MATCH (d:Director)\\nWHERE d.born < date('1950-...   \n",
       "2  MATCH (u:User)-[r:RATED]->(m:Movie)\\nWITH u, c...   \n",
       "3  MATCH (a:Actor)-[:ACTED_IN]->(m:Movie)\\nWITH a...   \n",
       "4  MATCH (m:Movie)-[:IN_GENRE]->(g:Genre)\\nWHERE ...   \n",
       "\n",
       "                       type         database  syntax_error  timeout  \\\n",
       "0  Simple Retrieval Queries  recommendations         False    False   \n",
       "1  Simple Retrieval Queries  recommendations         False    False   \n",
       "2  Simple Retrieval Queries  recommendations         False    False   \n",
       "3  Simple Retrieval Queries  recommendations         False    False   \n",
       "4  Simple Retrieval Queries  recommendations         False    False   \n",
       "\n",
       "   returns_results  \n",
       "0             True  \n",
       "1             True  \n",
       "2             True  \n",
       "3             True  \n",
       "4             True  "
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "final_df = combined_df[\n",
    "    [\n",
    "        \"question\",\n",
    "        \"cypher\",\n",
    "        \"type\",\n",
    "        \"database\",\n",
    "        \"syntax_error\",\n",
    "        \"timeout\",\n",
    "        \"returns_results\",\n",
    "    ]\n",
    "]\n",
    "final_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "96141532-be3d-4d17-bc03-c148048d6fd7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Distribution for syntax_error:\n",
      " syntax_error\n",
      "False    4772\n",
      "True      194\n",
      "Name: count, dtype: int64\n",
      "Distribution for timeout:\n",
      " timeout\n",
      "False    4898\n",
      "True       68\n",
      "Name: count, dtype: int64\n",
      "Distribution for returns_results:\n",
      " returns_results\n",
      "True     3987\n",
      "False     979\n",
      "Name: count, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "# Assume df is your DataFrame and col1, col2, col3 are the boolean columns\n",
    "distribution_col1 = final_df[\"syntax_error\"].value_counts()\n",
    "distribution_col2 = final_df[\"timeout\"].value_counts()\n",
    "distribution_col3 = final_df[\"returns_results\"].value_counts()\n",
    "\n",
    "print(\"Distribution for syntax_error:\\n\", distribution_col1)\n",
    "print(\"Distribution for timeout:\\n\", distribution_col2)\n",
    "print(\"Distribution for returns_results:\\n\", distribution_col3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "5317c5d2-add9-4deb-a990-5e2031a2b89c",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_df.to_csv(\"text2cypher_claudeopus.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cba3878-5d2e-4293-8392-f3ced39be392",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
